# -*-encoding: Utf-8 -*-
# @File : get_cookies.py
# author: HJT
# Time：2023/11/13
import requests
import warnings
import re

from utils.yaml_util import read_yaml

try:
    from encrypt import encrypt_password
except ImportError:
    from .encrypt import encrypt_password

cfg_info: dict = read_yaml()

BASE_URL = cfg_info.get('base_url')
USERNAME = cfg_info.get('login').get('user')
PASSWORD = cfg_info.get('login').get('password')


def login():
    """
    :param username: 登录用户名
    :param password: 密码
    :return: cookie实例
    """
    base_url = BASE_URL
    login_url = base_url + '/login/'
    username = USERNAME
    password = PASSWORD
    warnings.filterwarnings('ignore')
    # 公共配置
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36",
    }

    # 1. 请求1：get请求蓝鲸控制台页面
    session = requests.session()
    login_page = session.get(login_url, headers=headers, verify=False)

    # 请求1的响应提取：①提取登录页面html中，name="csrfmiddlewaretoken"; ②提取加密公钥PASSWORD_RSA_PUBLIC_KEY
    set_cookie = login_page.headers['Set-Cookie']
    pattern = r"bklogin_csrftoken=(.*?);"
    csrfmiddlewaretoken = re.search(pattern, set_cookie).group(1)
    pattern = r'PASSWORD_RSA_PUBLIC_KEY = "(.*?)"'
    PASSWORD_RSA_PUBLIC_KEY = re.search(pattern, login_page.text).group(1)
    # 2. 请求2：蓝鲸控制台登录
    # 密码加密处理
    password = encrypt_password(password, PASSWORD_RSA_PUBLIC_KEY)
    # 登录参数
    login_data = {
        "csrfmiddlewaretoken": csrfmiddlewaretoken,  # data中的重要参数，服务器会校验
        "username": username,
        "password": password,
        "next": "",
        "app_id": ""
    }
    # print(f'login_data数据：{login_data}')
    headers['Referer'] = login_url  # headers中的重要参数，服务器会校验
    login_resp = session.post(login_url, data=login_data, headers=headers, verify=False)
    # print(f'post方式的登录结果：{login_resp.status_code}')
    # print_log(f'请求2的响应：{login_resp.text}')
    # 断言响应结果：重定向结果的状态码==302
    assert login_resp.status_code < 400, f'{login_url}登录失败，请检查用户名、密码'

    # 3. 请求3：访问用户管理模块
    user_manage_url = base_url + '/o/bk_user_manage'
    # print_log(f'请求3的url：{url}')
    user_page = session.get(user_manage_url, headers=headers, verify=False)
    # print_log(f'请求3的响应：{itsm_page.text}')
    # 简单断言：响应是否成功
    assert user_page.status_code < 300, f'请求失败，URL:{user_manage_url}'

    # 获取用户身份信息(cookie)
    cookies = session.cookies.get_dict()
    if cookies is None:
        raise Exception("登录失败，cookie为空值")

    cookies_str = ""

    for k, v in cookies.items():
        row = f'{k}={v};'
        cookies_str += row

    set_cookie_header = user_page.headers.get('Set-Cookie')
    token = re.search('bkuser_csrftoken=(.*?);', set_cookie_header).group(1)
    print(f'登录成功！  登录用户：{username}， 密码：{password}')

    return cookies_str, token


if __name__ == '__main__':
    print(login())
